import torch
import torchvision
import torchvision.models as models
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.segmentation import deeplabv3_resnet101


# An instance of your model.
rcnn_model = fasterrcnn_resnet50_fpn(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True)
deeplab_v3 = deeplabv3_resnet101(pretrained=True, progress=True, num_classes=21)
model = torchvision.models.resnet34(pretrained=True)

# rcnn_model.eval()
# deeplab_v3.eval()
# model.eval()
# # An example input you would normally provide to your model's forward() method.
# example = torch.rand(1, 3, 224, 224)
#
# # Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
# traced_script_module = torch.jit.trace(model, example)
#
# traced_script_module.save("traced_resnet_model.pt")
